package algorithm;

import DataStructure.CharStack;
import DataStructure.Tree;
import DataStructure.TreeNode;
import it.unisa.dia.gas.jpbc.Element;
import it.unisa.dia.gas.jpbc.Field;
import it.unisa.dia.gas.jpbc.Pairing;
import util.PolyItem;
import util.Polynomial;
import util.Symbol;
import util.TaiHashUtil;

import java.math.BigInteger;
import java.nio.charset.StandardCharsets;
import java.util.*;

/**
 * CPABE函数的加密类，用户完成初始化之后，输入策略和明文即可进行加密
 * 更具策略的逆波兰表达式生成基于二叉树的访问树结构进行加密
 */
public class CPABEenc {
    private Pairing pairing;
    private CPABEPK pubKey;
    private Tree tree;
    private Element s;
    private BigInteger order;
    private int countAtt;
    private Element M;

    /**
     * 通过输入计算过程中用到的pairing和public Key来初始化函数
     * @param pairing 过程中用到的配对
     * @param pubKey 公钥类
     */

    public CPABEenc(Pairing pairing, CPABEPK pubKey, BigInteger Order) {
        this.pairing = pairing;
        this.pubKey = pubKey;
        countAtt=0;
        this.order=Order;
    }

    /**
     * CPABE加密文件
     * @param plaintext 明文
     * @param policy 策略，策略支持括号，#标识或，*标识与，*的优先级大于#
     * @return 加密过后的密文
     * @throws Exception
     */

    public CPABECT encrypt(String plaintext,String policy) throws Exception {
        parsePolicy(policy);
        //通过将输入的字符串映射到椭圆曲线群上，进行操作
        byte[] bText= plaintext.getBytes(StandardCharsets.UTF_8);
        BigInteger temp = new BigInteger(plaintext);
        //获取椭圆曲线基于的整数域
        Field zn=pairing.getZr();
        Field gt=pairing.getGT();
//        M=gt.newElement(temp).getImmutable();
        M=gt.newElementFromBytes(temp.toByteArray()).getImmutable();
        Element cm=pubKey.getEa().powZn(s.duplicate()).mul(M).getImmutable();

        Element C=pubKey.getH().powZn(s).getImmutable();
        EAUnion[] cy1=new EAUnion[countAtt];
        EAUnion[] cy2=new EAUnion[countAtt];
        //标记密文控件数组的位置
        int tag=0;
        List<TreeNode> allNode=tree.getAllNode();
        for (int i = 0; i < allNode.size(); i++) {
            TreeNode node=allNode.get(i);
            String attribute=node.getAttribute();
            if(attribute!=null){
                //叶节点才进行计算
                cy1[tag]=new EAUnion();
                Element pValue=zn.newElement(node.getPolynomial().calValue(BigInteger.ZERO)).getImmutable();
                cy1[tag].seteData(pubKey.getG().powZn(pValue).getImmutable());
                cy1[tag].setAttribute(attribute);

                byte[] hashAtt= TaiHashUtil.getMD5Str(attribute).getBytes(StandardCharsets.UTF_8);
                Element hashValue=pairing.getG1().newElementFromHash(hashAtt,0,hashAtt.length).getImmutable();
                cy2[tag]=new EAUnion();
                Element pValue2=zn.newElement(node.getPolynomial().calValue(BigInteger.ZERO)).getImmutable();
                cy2[tag].seteData(hashValue.powZn(pValue2).getImmutable());
                cy2[tag].setAttribute(attribute);
                tag++;
            }
            node.clearPoly();
        }

        return new CPABECT(tree,cm,C,cy1,cy2,policy);

    }



    /**
     * 构建策略树
     * @param policy 用户输入的策略
     * @throws Exception 策略输入格式有问题
     */
    public void parsePolicy(String policy) throws Exception {
        if (policy.length()==0){
            throw new Exception("策略不能为空/policy can`t be empty");
        }
        //初始化策略树
        List<String> rPolishType=reversePolishType(policy.toCharArray());
        String outData=listOutput(rPolishType);
        System.out.println("策略策略解析结果为："+outData);
        buildBinaryTree(rPolishType);
    }

    /**
     * 将用户输入的策略转变为逆波兰表达式，为建树做准备
     * 并且能够检测出策略格式是否有问题
     * 目前支持的检测点为：括号使用不当
     * @param expression 策略的字符数组，string数组的时候采用UTF-8编码
     * @return 逆波兰式中的每一项
     * @throws Exception 输入策略格式有问题
     */
    private List<String> reversePolishType(char[] expression) throws Exception {
        //TODO-pht:减少栈的使用量
        CharStack<Symbol> symbolStack = new CharStack<Symbol>();
        CharStack<Character> dataStack = new CharStack<Character>();
        List<String> output = new ArrayList<>();
        //output当中的的游标
        int tag = 0;
        //放入队列中的数据
        String inData;
        for (int i = 0; i < expression.length; i++) {
            Symbol temporary=null;
            //处理数据与运算符的关系
            switch (expression[i]) {
                case '(':
                    //左边一定是运算符，直接压入
                    Symbol leftBracket = new Symbol('(', getPriority('('));
                    symbolStack.push(leftBracket);
                    //左括号优先级最大，不需要做优先级检查
                    break;
                case ')':
                    //左边一定是数，将数据放入输出队列中
                    if (!dataStack.isEmpty()) {
                        inData = dataStack.getString();
                        output.add(tag++, inData);
                    }
                    //将运算符压入运算符栈中
                    Symbol rightBracket = new Symbol(')', getPriority(')'));
                    symbolStack.push(rightBracket);
                    //优先级检查
                    temporary=handlePriority(symbolStack);
                    while(temporary!=null){
                        inData=String.valueOf(temporary.getSym());
                        output.add(tag++,inData);
                        temporary=handlePriority(symbolStack);
                    }
                    break;
                case '*':
                    //压入运算符
                    Symbol starKey = new Symbol('*', getPriority('*'));
                    symbolStack.push(starKey);
                    if (!dataStack.isEmpty()) {
                        inData = dataStack.getString();
                        output.add(tag++, inData);
                    }
                    //优先级检查
                    temporary=handlePriority(symbolStack);
                    while(temporary!=null){
                        inData=String.valueOf(temporary.getSym());
                        output.add(tag++,inData);
                        temporary=handlePriority(symbolStack);
                    }
                    break;
                case '#':
                    //压入运算符
                    Symbol hashKey = new Symbol('#', getPriority('#'));
                    symbolStack.push(hashKey);
                    if(!dataStack.isEmpty()){
                        inData=dataStack.getString();
                        output.add(tag++,inData);
                    }
                    //优先级检查
                    temporary=handlePriority(symbolStack);
                    while(temporary!=null){
                        inData=String.valueOf(temporary.getSym());
                        output.add(tag++,inData);
                        temporary=handlePriority(symbolStack);
                    }
                break;
                default:
                    //将数据添加到数据栈中
                    countAtt++;
                    dataStack.push(expression[i]);
                break;
            }
        }
        //判断数据栈中是否为空，不为空则直接加入后续的输出当中
        if(!dataStack.isEmpty()){
            inData=dataStack.getString();
            output.add(tag++,inData);
        }
        //符号栈不为空，处理后续的东西
        //首先判断输入是否合法,同时生成符号的数组
        int stackSize=symbolStack.size();
        Symbol[] symbolList=new Symbol[stackSize];
        for (int i = 0; i < stackSize; i++) {
            symbolList[i]=symbolStack.pop();
            if(symbolList[i].getSym()==')'||symbolList[i].getSym()=='('){
                throw new  Exception("策略格式错误/wrong form of policy");
            }
        }
        //如果只有一个符号，则直接输出
        if(symbolList.length==1){
            output.add(tag++,String.valueOf(symbolList[0].getSym()));
        }else{
            //对符号数组按照优先级从高到底的顺序排列
            for (int i = 0; i <symbolList.length ; i++) {
                int maxtag=i;
                int maxPriority=symbolList[i].getPriority();
                for (int j = i+1; j < symbolList.length; j++) {
                    if(symbolList[j].getPriority()>maxPriority){
                        maxtag=j;
                        maxPriority=symbolList[j].getPriority();
                    }
                }
                Symbol temp=symbolList[i];
                symbolList[i]=symbolList[maxtag];
                symbolList[maxtag]=temp;
            }
            //将多出来的运算符添加到输出列表当中
            for (int i = 0; i < symbolList.length; i++) {
                inData=String.valueOf(symbolList[i].getSym());
                output.add(tag++,inData);
            }
        }
        return output;
    }

    private int getPriority(char symbol) {
        switch (symbol) {
            case '(':
                return 4;
            case '*':
                return 3;
            case '#':
                return 2;
            case ')':
                return 1;
            default:
                return 0;
        }
    }

    /**
     * 当符号入栈时进行优先级检查，并做弹出操作
     */
    private Symbol handlePriority(CharStack<Symbol> symbolStack){
        if(symbolStack.size()<=1){
            return null;
        }

        Symbol top=symbolStack.pop();
        Symbol down=symbolStack.pop();
        if(top.getPriority()>down.getPriority()){
            //当前插入的运算符的优先级大于前一个的时候，将两个符号按照顺序压回队列
            symbolStack.push(down);
            symbolStack.push(top);
        }else{
            //
            if(!(down.getSym()=='('||top.getSym()==')')){
                symbolStack.push(top);
                return down;
            }
            if(down.getSym()=='('&&top.getSym()!=')') {
                symbolStack.push(down);
                symbolStack.push(top);
            }
            if(down.getSym()!='('&&top.getSym()==')'){
                symbolStack.push(top);
                return down;
            }
            if(down.getSym()=='('&&top.getSym()==')'){
                return null;
            }
        }
        return null;
    }

    /**
     * 将输入的List中的各个字符串按照先后顺序排组合成字符串
     * @param listData 输入的list
     * @return 组合成的字符串
     */
    private String listOutput(List<String> listData){
                //组成输出字符串
        String outData="";
        for (int i = 0; i < listData.size(); i++) {
            outData=outData+listData.get(i);
        }
        return outData;
    }

    /**
     * 通过输入的逆波兰表达式的各个项，构建策略树
     * @param rPolishTye 逆波兰表达式的各个项
     */
    private void buildBinaryTree(List<String> rPolishTye){
        tree=new Tree();
        //如果只有一个节点的话
        if(rPolishTye.size()==1){
            TreeNode root=new TreeNode();
            //符号节点
            int[] child=new int[1];
            root.setChildren(child);
            root.setOrder(1);
            root.setThreshold(root.getOrder());
            //设置运算符子节点
            TreeNode onlyOne=new TreeNode(0,1,0,null);
            onlyOne.setChildren(null);
            onlyOne.setParent(1);
            onlyOne.setViewed(false);
            onlyOne.setAttribute(rPolishTye.get(0));
            root.setParent(-2);
            root.setNodeIndex(1);
            root.setAttribute(null);
            root.setViewed(false);
            tree.insertNode(onlyOne,onlyOne.getNodeIndex());
            tree.insertNode(root,root.getNodeIndex());

            //当只有一个属性节点的时候，根节点就是这个与运算节点
            tree.setRootIndex(1);
            fullFillTree();
            return;
        }
        //有多个节点的情况
        for (int i = 0; i < rPolishTye.size(); i++) {
            String data=rPolishTye.get(i);
            switch (getState(data)){
                case 1:
                    //符号节点
                    TreeNode nodeSymbol=new TreeNode();
                    int[] tuple=getNoParentList(tree,i);
                    nodeSymbol.setOrder(tuple[2]);
                    if(rPolishTye.get(i).equals("#")){
                        //或门，门限值为1
                        nodeSymbol.setThreshold(1);
                    }else{
                        //与门,门限值为阶的值
                        nodeSymbol.setThreshold(nodeSymbol.getOrder());
                    }
                    int[] children=new int[nodeSymbol.getOrder()];
                    //设置运算符子节点
                    for (int j = 0; j < children.length; j++) {
                        children[j]=tuple[j];
                    }
                    nodeSymbol.setChildren(children);
                    nodeSymbol.setParent(-2);
                    nodeSymbol.setNodeIndex(i);
                    nodeSymbol.setAttribute(null);
                    nodeSymbol.setViewed(false);
                    //修改子节点的父节点编号
                    int[] kid=nodeSymbol.getChildren();
                    for (int j = 0; j < nodeSymbol.getOrder(); j++) {
                        tree.getNode(kid[j]).setParent(nodeSymbol.getNodeIndex());
                    }
                    tree.insertNode(nodeSymbol,nodeSymbol.getNodeIndex());
                break;
                case 0:
                    //属性节点
                    TreeNode nodeData=new TreeNode();
                    nodeData.setOrder(0);
                    nodeData.setThreshold(1);
                    nodeData.setChildren(null);
                    //-2表示还没有设置父节点
                    nodeData.setParent(-2);
                    nodeData.setNodeIndex(i);
                    nodeData.setAttribute(rPolishTye.get(i));
                    nodeData.setViewed(false);
                    tree.insertNode(nodeData,nodeData.getNodeIndex());
                    break;
            }
        }
        //当有多个节点的时候，由逆波兰式的特点所决定，最后一个进入树的节点一定为根节点
        tree.setRootIndex(rPolishTye.size()-1);
        fullFillTree();
    }

    /**
     * 获得目前所有树节点中没有父节点的节点，用来建立父子节点之间的关系
     * @param tree 树
     * @param tag 目前节点的标号
     * @return int数组，size为3，第一个是右孩子节点的标号，第二个是左孩子节点的标号，
     */
    private int[] getNoParentList(Tree tree,int tag){
        //第一个值是左孩子节点的位置，第二个值是右孩子节点的位置，第三个值是有几个孩子
        int[] tuple=new int[3];
        int size=0;
        int index=0;
        for (int i = tag-1; i >=0; i--) {
            if(tree.getNode(i).getParent()==-2){
                //没有设置父节点
                tuple[index]=i;
                index++;
                size++;
                if(size==2){
                    break;
                }
            }
        }
        tuple[2]=size;
        return  tuple;
    }

    /**
     * 获取上一个节点的状态值
     * @param part 状态
     * @return 状态值
     */
    private int getState(String part){
        switch (part){
            case "(":
                return 1;
            case ")":
                return 1;
            case "#":
                return 1;
            case "*":
                return 1;
            default:
                return 0;
        }
    }

    /**
     * 在树由超过两个节点构成的时候，填充每一个节点在CPABE算法种设定的值
     * 该函数调用另外一个递归函数来完成该操作。
     */
    private void fullFillTree(){
        //TODO-pht:将填充多项式用的树改为Element类型

        //填充父节点的值
        TreeNode root=tree.getNode(tree.getRootIndex());
        Random random=new Random();
        BigInteger sTemp=new BigInteger(160,random);
        sTemp=sTemp.mod(order);
        s=pairing.getZr().newElement(sTemp).getImmutable();
        Polynomial p=genneratePoly(root.getThreshold()-1,sTemp);
        root.setPolynomial(p);
        root.setIndex(1);
        //填充余下所有节点的值
        int[] children=root.getChildren();
        if(children.length==2){
            fullNode(tree.getNode(children[0]),1);
            fullNode(tree.getNode(children[1]),2);
        }else if(children.length==1){
            fullNode(tree.getNode(children[0]),1);
        }
    }

    /**
     * 递归填充一个节点的子节点
     * @param node
     */
    private void fullNode(TreeNode node,int index){
        //TODO-pht:修改为非递归版本
        node.setIndex(index);
//        BigInteger constant=BigInteger.valueOf(tree.getNode(node.getParent()).getIndex());
        BigInteger constant=tree.getNode(node.getParent())
                .getPolynomial()
                .calValue(BigInteger.valueOf(node.getIndex()));
//        constant=constant.mod(order);
        Polynomial p=genneratePoly(node.getThreshold()-1,constant);
        node.setPolynomial(p);
        //当没有子节点的时候退出递归
        if(node.getChildren()==null){
            return;
        }
        int[] childre=node.getChildren();
        if(childre.length==2){
            fullNode(tree.getNode(childre[0]),0);
            fullNode(tree.getNode(childre[1]),1);
        }else if(childre.length==1){
            fullNode(tree.getNode(childre[0]),0);
        }
    }

    /**
     * 生成多项式
     * @param order 多项式的阶
     * @param constant 0次项的系数
     * @return 多项式
     */
    private Polynomial genneratePoly(int order,BigInteger constant){
        PolyItem[] items=new PolyItem[order+1];
        for (int i = 0; i < order; i++) {
            Random random=new Random();
            BigInteger coe=new BigInteger(160,random);
            int itemOrder=order-i;
            PolyItem newItem=new PolyItem(coe,itemOrder);
            items[i]=newItem;
        }
        PolyItem last=new PolyItem(constant,0);
        items[order]=last;
        Polynomial polynomial=new Polynomial(items);
        return polynomial;
    }
    public Element getM(){return M;}
}
